###***Model for viscous spreading.  Can calculate evolution of the surface mass density for either constant or variable viscosity.  Inputs
###***are x = switch for constant or variable (0 cont, 1 var), y is the timestep, dt the time between impacts, and z is the location of the interior bin
from __future__ import division
from math import *
from Init_Cond import *
Snew = []
Tnew = []

def f(x, M, t):        #X is a switch between constant viscosity and variable
    del Snew [:]
    del S[:]
    del nu[:]
    del Tnew[:]

    for a in range(int(N)):
        if x == 0:
            nu.append(1000.0*(r[a]/R_Mars)**2.0)
        else:
            if sigma[a] <= 1.0e-6:
                nu.append(0.0)
                S.append(sigma[a]*X[a])
                continue

            r_hstar = R_Embryo/(2.0*r_pdisk)*(2.0*m_pdisk/(3.0*M))**(1.0/3.0)
            tau = pi*r_pdisk**2*sigma[a]/m_pdisk


            if r_hstar < 1.0:
                sigma_r = 0.5*(1+tanh((2.*r_hstar-1)/(r_hstar*(r_hstar-1))))*sqrt(G*m_pdisk/r_pdisk) +  0.5*(1-tanh((2.*r_hstar-1)/(r_hstar*(r_hstar-1))))*(2.0*r_pdisk*w[a])
            else:
                sigma_r = sqrt(G*m_pdisk/r_pdisk)


            Q = w[a]*sigma_r/(3.36*G*(sigma[a]+0.000001))

            if Q <= 4.0:

                nu_trans_stable = sigma_r**2.0/(2.0*w[a])*(0.46*tau/(1.0+tau**2.0))
                nu_grav_stable = 0.0
                nu_trans_unstable = 13.0*r_hstar**5.0*G**2.0*sigma[a]**2.0/w[a]**3.0
                nu_grav_unstable = nu_trans_unstable

                y = Q/4.
                nu_trans = 0.5*(1+tanh((2.*y-1)/(y*(y-1))))*(nu_trans_stable) +  0.5*(1-tanh((2.*y-1)/(y*(y-1))))*(nu_trans_unstable)
                nu_grav = 0.5*(1+tanh((2.*y-1)/(y*(y-1))))*(nu_grav_stable) +  0.5*(1-tanh((2.*y-1)/(y*(y-1))))*(nu_grav_unstable)

            else:
                nu_trans = sigma_r**2.0/(2.0*w[a])*(0.46*tau/(1.0+tau**2.0))
                nu_grav = 0.0

            nu_coll = r_pdisk**2.0*w[a]*tau
            nu.append(nu_trans + nu_grav + nu_coll)

        S.append(sigma[a]*X[a])

def g(inside, deltat, M):

#solution approach from bath pringle '81
    for a in range(int(N)):
        if a < inside:  #inside surface of embryo
            Snew.append(S[a])
            Tnew.append(0.0)    #torque from satellite onto disk
        elif a == inside:   #at first disk bin

            Snew.append(S[a] + 12.*deltat/(X[a]**2.*deltaX**2.)*(S[a+1]*nu[a+1] - 2.*S[a]*nu[a]))   #boundary condition for inner bin
            Tnew.append(deltat/(r[a]*pi*sqrt(G*M))*(sqrt(r[a+1])*Torque_to_disk[a+1] - sqrt(r[a])*Torque_to_disk[a])/deltar)  #boundary condition

        elif a == int(N) - 1:   #at last disk bin
            Snew.append(S[a])
            Tnew.append(0.0)
        else:   #rest of the disk
            Snew.append(S[a] + 12.*deltat/(X[a]**2.*deltaX**2.)*(S[a-1]*nu[a-1] + S[a+1]*nu[a+1] - 2.*S[a]*nu[a]))
            Tnew.append(deltat/(r[a]*pi*sqrt(G*M))*(sqrt(r[a+1])*Torque_to_disk[a+1] - sqrt(r[a-1])*Torque_to_disk[a-1])/(2.*deltar))

        r_L = 2**(2./3.)*r[a]   #find farthest satellite bin this disk bin is in resonance with
        i = int(round((r_L - r_I)/deltar - 0.5))    #identify satellite bin
        if i >= int(N): #if it's outside disk move one
            continue


def o(Mars_Accrete, inside):  #input is Mars_Accrete, so is output


    if (Snew[inside]/X[inside] - Tnew[inside])*deltaA[inside] < m[inside] and (Snew[inside]/X[inside] - Tnew[inside]) >= 0.0:    #If first mass carrying bin has decreased in mass, that mass has passed to the interior and adds to Mars_Accrete
        #####******~~~~~  THIS IS NOT IMPLEMENTED WELL
        # return Mars_Accrete + (S[inside]/X[inside]-Snew[inside]/X[inside])*deltaA[inside]
        return Mars_Accrete + (S[inside]/X[inside]-(Snew[inside]/X[inside] - Tnew[inside]))*deltaA[inside]
    else:
        return Mars_Accrete

def s():

    for a in range(int(N)): #Set update sigma and mass of each bin to the new values

        sigma[a] = Snew[a]/X[a] - Tnew[a]

        m[a] = sigma[a]*deltaA[a]
        Torque_to_disk[a] = 0.0